import os
import sys

from openai import OpenAI
sys.path.append('/JiuTian/swarm/swarm-main/swarm')
from swarm import Swarm,Agent


#测试基本功能
client = OpenAI(api_key=os.environ["DEEPSEEK_API_KEY"], base_url=os.environ["DEEPSEEK_BASE_URL"])
# client = OpenAI(api_key="sk-dmtzbscxyyaijoprsokdsxcjxazjrvrgzxhinvkwgomkonhg", base_url="https://api.siliconflow.cn/v1")

def transfer_to_agent_b():
    return agent_b

def transfer_to_agent_c():
    return agent_c

def transfer_to_agent_d():
    return agent_d

# model_name = "deepseek-ai/DeepSeek-V3"
model_name = "deepseek-chat"
#测试多智能体功能
swarm_client = Swarm(client)
agent_a = Agent(name = "Agent A",instructions= "你是一个乐于助人的助手",functions=[transfer_to_agent_b,transfer_to_agent_c,transfer_to_agent_d],model=model_name)
agent_b = Agent(name = "Agent B",instructions= "只使用俳句回答。",model=model_name)
agent_c = Agent(name = "Agent C",instructions= "只使用儿歌回答。",model=model_name)
agent_d = Agent(name = "Agent D",instructions= "只使用英语回答。",model=model_name)

response = swarm_client.run(agent=agent_a,messages=[{"role":"user","content":"我想和智能体B谈话，请问你你喜欢红楼梦吗？"}])
print(response.messages[-1]["content"])

response = swarm_client.run(agent=agent_a,messages=[{"role":"user","content":"我想和智能体C谈话，请问你你喜欢红楼梦吗？"}])
print(response.messages[-1]["content"])

response = swarm_client.run(agent=agent_a,messages=[{"role":"user","content":"我想和智能体D谈话，请问你你喜欢红楼梦吗？"}])
print(response.messages[-1]["content"])